Thanks a lot Philip, LGTM.

I have 2 follow-up questions / comments, which we can address (or not) in other PRs:

• As discussed offline, JIT "isn't supported" for the datapoints on these dispatchers. What we mean by "not supported" is that the datapoints will be treated as a pure Tensors, and the output will be a Tensor too. I assume this can lead to unexpected results and confusion, e.g.:

  mask = Mask(...)
  out = T(mask)
  # - unaware user expects a Mask but gets a Tensor
  # - `T` treats mask as an image, not as a mask!
  #    Might result in different behaviour (very bad) or confusing error message?

  so: should we fail loudly with a good error message in the disaptchers when is_scripting() is True and is_simple_tensor() is False?

• Side note: should we be more lenient with mypy checks in the prototype area?
  We introduced strict checks in make mypy more strict for prototype datasets #4513 (and I approved the PR, so I'll concede I'm sort of rebutting myself here).
  But the the prototype area is where we want to be able to move fast, and try potentially crazy things without having to worry too much about type checks / docs / super clean tests. In this specific case, the introduction of warn_return_any could have been the reason for these convoluted checks in the first place which complicates the code and slows down iteration speed. Perhaps we may want to relax the mypy requirements in order to experiment faster in the prototype area?